package com.atguigu.hadoop.examples.mapreduce.terasort;


import java.io.DataInputStream;
import java.io.IOException;
import java.io.PrintStream;
import java.net.URI;
import org.apache.hadoop.conf.Configurable;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.conf.Configured;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.JobContext;
import org.apache.hadoop.mapreduce.Partitioner;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
import org.apache.hadoop.util.Tool;
import org.apache.hadoop.util.ToolRunner;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class TeraSort extends Configured
        implements Tool
{
    private static final Logger LOG = LoggerFactory.getLogger(TeraSort.class);

    public static boolean getUseSimplePartitioner(JobContext job)
    {
        return job.getConfiguration().getBoolean(TeraSortConfigKeys.USE_SIMPLE_PARTITIONER
                .key(), false);
    }

    public static void setUseSimplePartitioner(Job job, boolean value)
    {
        job.getConfiguration().setBoolean(TeraSortConfigKeys.USE_SIMPLE_PARTITIONER
                .key(), value);
    }

    public static int getOutputReplication(JobContext job) {
        return job.getConfiguration().getInt(TeraSortConfigKeys.OUTPUT_REPLICATION
                .key(), 1);
    }

    public static void setOutputReplication(Job job, int value)
    {
        job.getConfiguration().setInt(TeraSortConfigKeys.OUTPUT_REPLICATION.key(), value);
    }

    private static void usage() throws IOException
    {
        System.err.println("Usage: terasort [-Dproperty=value] <in> <out>");
        System.err.println("TeraSort configurations are:");
        for (TeraSortConfigKeys teraSortConfigKeys : TeraSortConfigKeys.values()) {
            System.err.println(teraSortConfigKeys.toString());
        }
        System.err.println("If you want to store the output data as erasure code striping file, just make sure that the parent dir of <out> has erasure code policy set");
    }

    public int run(String[] args)
            throws Exception
    {
        if (args.length != 2) {
            usage();
            return 2;
        }
        LOG.info("starting");
        Job job = Job.getInstance(getConf());
        Path inputDir = new Path(args[0]);
        Path outputDir = new Path(args[1]);
        boolean useSimplePartitioner = getUseSimplePartitioner(job);
        TeraInputFormat.setInputPaths(job, new Path[] { inputDir });
        FileOutputFormat.setOutputPath(job, outputDir);
        job.setJobName("TeraSort");
        job.setJarByClass(TeraSort.class);
        job.setOutputKeyClass(Text.class);
        job.setOutputValueClass(Text.class);
        job.setInputFormatClass(TeraInputFormat.class);
        job.setOutputFormatClass(TeraOutputFormat.class);
        if (useSimplePartitioner) {
            job.setPartitionerClass(SimplePartitioner.class);
        } else {
            long start = System.currentTimeMillis();
            Path partitionFile = new Path(outputDir, "_partition.lst");

            URI partitionUri = new URI(partitionFile.toString() + "#" + "_partition.lst");
            try
            {
                TeraInputFormat.writePartitionFile(job, partitionFile);
            } catch (Throwable e) {
                LOG.error(e.getMessage());
                return -1;
            }
            job.addCacheFile(partitionUri);
            long end = System.currentTimeMillis();
            System.out.println("Spent " + (end - start) + "ms computing partitions.");
            job.setPartitionerClass(TotalOrderPartitioner.class);
        }

        job.getConfiguration().setInt("dfs.replication", getOutputReplication(job));
        int ret = job.waitForCompletion(true) ? 0 : 1;
        LOG.info("done");
        return ret;
    }

    public static void main(String[] args)
            throws Exception
    {
        int res = ToolRunner.run(new Configuration(), new TeraSort(), args);
        System.exit(res);
    }

    public static class SimplePartitioner extends Partitioner<Text, Text>
            implements Configurable
    {
        int prefixesPerReduce;
        private static final int PREFIX_LENGTH = 3;
        private Configuration conf = null;

        public void setConf(Configuration conf) { this.conf = conf;
            this.prefixesPerReduce = ((int)Math.ceil(16777216.0F / conf
                    .getInt("mapreduce.job.reduces", 1)));
        }

        public Configuration getConf()
        {
            return this.conf;
        }

        public int getPartition(Text key, Text value, int numPartitions)
        {
            byte[] bytes = key.getBytes();
            int len = Math.min(3, key.getLength());
            int prefix = 0;
            for (int i = 0; i < len; i++) {
                prefix = prefix << 8 | 0xFF & bytes[i];
            }
            return prefix / this.prefixesPerReduce;
        }
    }

    static class TotalOrderPartitioner extends Partitioner<Text, Text>
            implements Configurable
    {
        private TrieNode trie;
        private Text[] splitPoints;
        private Configuration conf;

        private static Text[] readPartitions(FileSystem fs, Path p, Configuration conf)
                throws IOException
        {
            int reduces = conf.getInt("mapreduce.job.reduces", 1);
            Text[] result = new Text[reduces - 1];
            DataInputStream reader = fs.open(p);
            for (int i = 0; i < reduces - 1; i++) {
                result[i] = new Text();
                result[i].readFields(reader);
            }
            reader.close();
            return result;
        }

        private static TrieNode buildTrie(Text[] splits, int lower, int upper, Text prefix, int maxDepth)
        {
            int depth = prefix.getLength();
            if ((depth >= maxDepth) || (lower == upper)) {
                return new LeafTrieNode(depth, splits, lower, upper);
            }
            InnerTrieNode result = new InnerTrieNode(depth);
            Text trial = new Text(prefix);

            trial.append(new byte[1], 0, 1);
            int currentBound = lower;
            for (int ch = 0; ch < 255; ch++) {
                trial.getBytes()[depth] = ((byte)(ch + 1));
                lower = currentBound;
                while ((currentBound < upper) &&
                        (splits[currentBound].compareTo(trial) < 0))
                {
                    currentBound++;
                }
                trial.getBytes()[depth] = ((byte)ch);
                result.child[ch] = buildTrie(splits, lower, currentBound, trial, maxDepth);
            }

            trial.getBytes()[depth] = -1;
            result.child['ÿ'] = buildTrie(splits, currentBound, upper, trial, maxDepth);

            return result;
        }

        public void setConf(Configuration conf) {
            try {
                FileSystem fs = FileSystem.getLocal(conf);
                this.conf = conf;
                Path partFile = new Path("_partition.lst");
                this.splitPoints = readPartitions(fs, partFile, conf);
                this.trie = buildTrie(this.splitPoints, 0, this.splitPoints.length, new Text(), 2);
            } catch (IOException ie) {
                throw new IllegalArgumentException("can't read partitions file", ie);
            }
        }

        public Configuration getConf() {
            return this.conf;
        }

        public int getPartition(Text key, Text value, int numPartitions)
        {
            return this.trie.findPartition(key);
        }

        static class LeafTrieNode extends TeraSort.TotalOrderPartitioner.TrieNode
        {
            int lower;
            int upper;
            Text[] splitPoints;

            LeafTrieNode(int level, Text[] splitPoints, int lower, int upper)
            {
                super(level);
                this.splitPoints = splitPoints;
                this.lower = lower;
                this.upper = upper;
            }
            int findPartition(Text key) {
                for (int i = this.lower; i < this.upper; i++) {
                    if (this.splitPoints[i].compareTo(key) > 0) {
                        return i;
                    }
                }
                return this.upper;
            }
            void print(PrintStream strm) throws IOException {
                for (int i = 0; i < 2 * getLevel(); i++) {
                    strm.print(' ');
                }
                strm.print(this.lower);
                strm.print(", ");
                strm.println(this.upper);
            }
        }

        static class InnerTrieNode extends TeraSort.TotalOrderPartitioner.TrieNode
        {
            private TeraSort.TotalOrderPartitioner.TrieNode[] child = new TeraSort.TotalOrderPartitioner.TrieNode[256];

            InnerTrieNode(int level) {
                super(level);
            }
            int findPartition(Text key) {
                int level = getLevel();
                if (key.getLength() <= level) {
                    return this.child[0].findPartition(key);
                }
                return this.child[(key.getBytes()[level] & 0xFF)].findPartition(key);
            }
            void setChild(int idx, TeraSort.TotalOrderPartitioner.TrieNode child) {
                this.child[idx] = child;
            }
            void print(PrintStream strm) throws IOException {
                for (int ch = 0; ch < 256; ch++) {
                    for (int i = 0; i < 2 * getLevel(); i++) {
                        strm.print(' ');
                    }
                    strm.print(ch);
                    strm.println(" ->");
                    if (this.child[ch] != null)
                        this.child[ch].print(strm);
                }
            }
        }

        static abstract class TrieNode
        {
            private int level;

            TrieNode(int level)
            {
                this.level = level; }
            abstract int findPartition(Text paramText);

            abstract void print(PrintStream paramPrintStream) throws IOException;

            int getLevel() { return this.level; }

        }
    }
}